import requests
import json
from openai import OpenAI
from llm.qianwen_prompt import *

qwen_2p5_72b_url= 'https://jinan21207.interior.aixcoder.com' #'http://10.103.255.12:21207'
headers ={
    'Content-Type': 'application/json',
    'Authorization': 'Bearer token-test-jsy',
}

def qianwen_stream(text:str="你好"):
    try:
        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            # api_key=os.getenv("DASHSCOPE_API_KEY"), 
            api_key="sk-3f966520a2684aea9367de0ebe63508e",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        completion = client.chat.completions.create(
            model="qwen2.5-72b-instruct", # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
            messages=[
                {
                    'role': 'system', 
                    'content': qianwen_system_prompt
                },{
                    'role': 'user', 
                    'content': qianwen_default_prompt.replace("{{ input_psuedo_code }}",text)
                }],
            temperature=0.7,
            top_p=0,
            stream=True,
            stream_options={"include_usage": True}
            )
        res = ""
        for chunk in completion:
            res += chunk.choices[0].delta.content
            print(res)
            print("--"*30)
            yield chunk.choices[0].delta.content
    except Exception as e:
        print(e)
    '''
    data ={
        'model': 'aix-76b-llm',
        'messages': [
            {
                'role': 'system',
                'content': qianwen_system_prompt
            },
            {
                'role': 'user',
                'content': qianwen_default_prompt.replace("{{ 替换输入内容 }}",text)
            }
        ],
        'temperature': 0.0,
        'stream': True,
    }
    logging.info(f"request data:{data}")
    try:
        res=""
        with requests.post(qwen_2p5_72b_url+'/v1/chat/completions',headers=headers,json=data,stream=True) as r:
            r.raise_for_status()
            for line in r.iter_lines():
                str_line = line.decode('utf-8')
                # yield str_line
                if str_line:
                    str_line=str_line[6:]
                    if str_line == '[DONE]':
                        print("DONE")
                    else:
                        json_line=json.loads(str_line)["choices"][0]["delta"]
                        if 'content' in json_line:
                            res= res+json_line['content']
                            print("res:",res)
                            yield res

    except Exception as e:
        logging.error(e)
        yield f"data: Error occurred: {str(e)}\n\n"
    '''

def qianwen(text:str="你好"):
    try:
        client = OpenAI(
            api_key="sk-3f966520a2684aea9367de0ebe63508e",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        completion = client.chat.completions.create(
            model="qwen2.5-72b-instruct",
            messages=[
                {
                    'role': 'system', 
                    'content': qianwen_system_prompt
                },
                {
                    'role': 'user', 
                    'content': qianwen_default_prompt.replace("{{ input_psuedo_code }}",text)
                }
            ],
            temperature=0.7,
            top_p=0,
        )
        response_data = json.loads(completion.model_dump_json())
        content = response_data["choices"][0]["message"]["content"]
        return get_pseudo_code_from_model_res(content=content)
    except Exception as e:
        print(e)

def call_model(text: str, url: str, model: str):
     payload = {
          "model": model,
          "prompt": qianwen_default_prompt.replace("{{ input_psuedo_code }}",text),
          "stream": False,
    }
     try:
          response = requests.post(url=url, json=payload)
          response.raise_for_status()
          content = json.loads(response.text)
          return get_pseudo_code_from_model_res(content)
     except Exception as e:
          print(e)

def get_pseudo_code_from_model_res(content):
    res = content[:content.rfind("```")]
    res = res[(res.rfind("```") + 3):]
    res = res[(res.find("\n") + 1):]
    return res